use std::{collections::HashMap, path::PathBuf, sync::Arc};

use baml_runtime::{BamlRuntime, TripWire};
use baml_types::BamlValue;
use function_result::FunctionResult;
use function_result_stream::FunctionResultStream;
use internal_baml_core::feature_flags::FeatureFlags;
use magnus::{class, function, method, prelude::*, Error, RArray, RHash, RModule, Ruby};
use types::{
    log_collector::Collector, request::HTTPRequest, runtime_ctx_manager::RuntimeContextManager,
};

mod function_result;
mod function_result_stream;
mod ruby_to_json;
mod types;

type Result<T> = std::result::Result<T, magnus::Error>;

// must be kept in sync with rb.define_class in the init() fn
#[magnus::wrap(class = "Baml::Ffi::BamlRuntime", free_immediately, size)]
struct BamlRuntimeFfi {
    inner: Arc<BamlRuntime>,
    t: Arc<tokio::runtime::Runtime>,
}

impl Drop for BamlRuntimeFfi {
    fn drop(&mut self) {
        use baml_runtime::runtime_interface::ExperimentalTracingInterface;
        match self.inner.flush() {
            Ok(_) => log::trace!("Flushed BAML log events"),
            Err(e) => log::error!("Error while flushing BAML log events: {e:?}"),
        }
    }
}

impl BamlRuntimeFfi {
    fn make_tokio_runtime(ruby: &Ruby) -> Result<tokio::runtime::Runtime> {
        // NB: libruby will panic if called from a non-Ruby thread, so we stick to the current thread
        // to avoid causing issues
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .map_err(|e| {
                Error::new(
                    ruby.exception_runtime_error(),
                    format!("Failed to start tokio runtime because:\n{e:?}"),
                )
            })
    }

    pub fn from_directory(
        ruby: &Ruby,
        directory: PathBuf,
        env_vars: HashMap<String, String>,
    ) -> Result<BamlRuntimeFfi> {
        let baml_runtime =
            match BamlRuntime::from_directory(&directory, env_vars, FeatureFlags::new()) {
                Ok(br) => br,
                Err(e) => {
                    return Err(Error::new(
                        ruby.exception_runtime_error(),
                        format!("{:?}", e.context("Failed to initialize BAML runtime")),
                    ))
                }
            };

        let rt = BamlRuntimeFfi {
            inner: Arc::new(baml_runtime),
            t: Arc::new(Self::make_tokio_runtime(ruby)?),
        };

        Ok(rt)
    }

    pub fn from_files(
        ruby: &Ruby,
        root_path: String,
        files: HashMap<String, String>,
        env_vars: HashMap<String, String>,
    ) -> Result<Self> {
        let baml_runtime =
            match BamlRuntime::from_file_content(&root_path, &files, env_vars, FeatureFlags::new())
            {
                Ok(br) => br,
                Err(e) => {
                    return Err(Error::new(
                        ruby.exception_runtime_error(),
                        format!("{:?}", e.context("Failed to initialize BAML runtime")),
                    ))
                }
            };

        let rt = BamlRuntimeFfi {
            inner: Arc::new(baml_runtime),
            t: Arc::new(Self::make_tokio_runtime(ruby)?),
        };

        Ok(rt)
    }

    pub fn create_context_manager(&self) -> RuntimeContextManager {
        RuntimeContextManager {
            inner: self
                .inner
                .create_ctx_manager(BamlValue::String("ruby".to_string()), None),
        }
    }

    pub fn call_function(
        ruby: &Ruby,
        rb_self: &BamlRuntimeFfi,
        function_name: String,
        args: RHash,
        ctx: &RuntimeContextManager,
        type_registry: Option<&types::type_builder::TypeBuilder>,
        client_registry: Option<&types::client_registry::ClientRegistry>,
        collector: RArray,
        env_vars: HashMap<String, String>,
    ) -> Result<FunctionResult> {
        let args = match ruby_to_json::RubyToJson::convert_hash_to_json(args) {
            Ok(args) => args.into_iter().collect(),
            Err(e) => {
                return Err(Error::new(
                    ruby.exception_syntax_error(),
                    format!("error while parsing call_function args:\n{e}"),
                ));
            }
        };

        let mut collectors = Vec::new();
        for i in collector.into_iter() {
            collectors.push(<&Collector>::try_convert(i)?.inner.clone());
        }

        let retval = match rb_self.t.block_on(rb_self.inner.call_function(
            function_name.clone(),
            &args,
            &ctx.inner,
            type_registry.map(|t| &t.inner),
            client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
            Some(collectors),
            env_vars,
            None,
            TripWire::new(None),
        )) {
            (Ok(res), _) => Ok(FunctionResult::new(res)),
            (Err(e), _) => Err(Error::new(
                ruby.exception_runtime_error(),
                format!(
                    "{:?}",
                    e.context(format!("error while calling {function_name}"))
                ),
            )),
        };

        retval
    }

    fn stream_function(
        ruby: &Ruby,
        rb_self: &BamlRuntimeFfi,
        function_name: String,
        args: RHash,
        ctx: &RuntimeContextManager,
        type_registry: Option<&types::type_builder::TypeBuilder>,
        client_registry: Option<&types::client_registry::ClientRegistry>,
        collector: RArray,
        env_vars: HashMap<String, String>,
    ) -> Result<FunctionResultStream> {
        let args = match ruby_to_json::RubyToJson::convert_hash_to_json(args) {
            Ok(args) => args.into_iter().collect(),
            Err(e) => {
                return Err(Error::new(
                    ruby.exception_syntax_error(),
                    format!("error while parsing stream_function args:\n{e}"),
                ));
            }
        };

        let mut collectors = Vec::new();
        for i in collector.into_iter() {
            collectors.push(<&Collector>::try_convert(i)?.inner.clone());
        }

        log::debug!("Streaming {function_name} with:\nargs: {args:#?}\nctx ???");

        let retval = match rb_self.inner.stream_function(
            function_name.clone(),
            &args,
            &ctx.inner,
            type_registry.map(|t| &t.inner),
            client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
            Some(collectors),
            env_vars,
            TripWire::new(None),
            None,
        ) {
            Ok(res) => Ok(FunctionResultStream::new(res, rb_self.t.clone())),
            Err(e) => Err(Error::new(
                ruby.exception_runtime_error(),
                format!(
                    "{:?}",
                    e.context(format!("error while calling {function_name}"))
                ),
            )),
        };

        retval
    }

    pub fn build_request(
        ruby: &Ruby,
        rb_self: &BamlRuntimeFfi,
        function_name: String,
        args: RHash,
        ctx: &RuntimeContextManager,
        type_registry: Option<&types::type_builder::TypeBuilder>,
        client_registry: Option<&types::client_registry::ClientRegistry>,
        env_vars: HashMap<String, String>,
        stream: bool,
    ) -> Result<HTTPRequest> {
        let args = match ruby_to_json::RubyToJson::convert_hash_to_json(args) {
            Ok(args) => args.into_iter().collect(),
            Err(e) => {
                return Err(Error::new(
                    ruby.exception_syntax_error(),
                    format!("error while parsing call_function args:\n{e}"),
                ));
            }
        };

        rb_self
            .inner
            .build_request_sync(
                function_name.clone(),
                &args,
                &ctx.inner,
                type_registry.map(|t| &t.inner),
                client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
                stream,
                env_vars,
            )
            .map(HTTPRequest::from)
            .map_err(|e| {
                Error::new(
                    ruby.exception_runtime_error(),
                    format!(
                        "{:?}",
                        e.context(format!(
                            "error while building HTTP request for function {function_name}"
                        ))
                    ),
                )
            })
    }

    pub fn parse_llm_response(
        ruby: &Ruby,
        rb_self: &BamlRuntimeFfi,
        function_name: String,
        llm_response: String,
        types: RModule,
        partial_types: RModule,
        allow_partials: bool,
        ctx: &RuntimeContextManager,
        type_registry: Option<&types::type_builder::TypeBuilder>,
        client_registry: Option<&types::client_registry::ClientRegistry>,
        env_vars: HashMap<String, String>,
    ) -> Result<magnus::Value> {
        let parsed = rb_self
            .inner
            .parse_llm_response(
                function_name.clone(),
                llm_response,
                allow_partials,
                &ctx.inner,
                type_registry.map(|t| &t.inner),
                client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
                env_vars,
            )
            .map_err(|e| {
                Error::new(
                    ruby.exception_runtime_error(),
                    format!(
                        "{:?}",
                        e.context(format!(
                            "error while parsing LLM response for function {function_name}"
                        ))
                    ),
                )
            })?;

        ruby_to_json::RubyToJson::serialize_baml(ruby, types, partial_types, allow_partials, parsed)
            .map_err(|e| {
                magnus::Error::new(
                    ruby.exception_type_error(),
                    format!("failed coercing BAML value to Ruby value: {e:?}"),
                )
            })
    }
}

fn invoke_runtime_cli(ruby: &Ruby, argv0: String, argv: Vec<String>) -> Result<u32> {
    match baml_cli::run_cli(
        std::iter::once(argv0).chain(argv).collect(),
        baml_runtime::RuntimeCliDefaults {
            output_type: baml_types::GeneratorOutputType::RubySorbet,
        },
    ) {
        Ok(exit_code) => Ok(exit_code.into()),
        Err(e) => Err(Error::new(
            ruby.exception_runtime_error(),
            format!(
                "{:?}",
                e.context("error while invoking baml-cli".to_string())
            ),
        )),
    }
}

#[magnus::init(name = "ruby_ffi")]
fn init(ruby: &Ruby) -> Result<()> {
    baml_log::init().map_err(|e| {
        Error::new(
            ruby.exception_runtime_error(),
            format!("Failed to initialize BAML logger: {e:#}"),
        )
    })?;

    let module = ruby.define_module("Baml")?.define_module("Ffi")?;

    module.define_module_function("invoke_runtime_cli", function!(invoke_runtime_cli, 2))?;

    // must be kept in sync with the magnus::wrap annotation
    let runtime_class = module.define_class("BamlRuntime", class::object())?;
    runtime_class.define_singleton_method(
        "from_directory",
        function!(BamlRuntimeFfi::from_directory, 2),
    )?;
    runtime_class
        .define_singleton_method("from_files", function!(BamlRuntimeFfi::from_files, 3))?;
    runtime_class.define_method(
        "create_context_manager",
        method!(BamlRuntimeFfi::create_context_manager, 0),
    )?;
    runtime_class.define_method("call_function", method!(BamlRuntimeFfi::call_function, 7))?;
    runtime_class.define_method(
        "stream_function",
        method!(BamlRuntimeFfi::stream_function, 7),
    )?;
    runtime_class.define_method("build_request", method!(BamlRuntimeFfi::build_request, 7))?;
    runtime_class.define_method(
        "parse_llm_response",
        method!(BamlRuntimeFfi::parse_llm_response, 9),
    )?;

    FunctionResult::define_in_ruby(&module)?;
    FunctionResultStream::define_in_ruby(&module)?;

    RuntimeContextManager::define_in_ruby(&module)?;

    types::type_builder::TypeBuilder::define_in_ruby(&module)?;
    types::type_builder::EnumBuilder::define_in_ruby(&module)?;
    types::type_builder::ClassBuilder::define_in_ruby(&module)?;
    types::type_builder::EnumValueBuilder::define_in_ruby(&module)?;
    types::type_builder::ClassPropertyBuilder::define_in_ruby(&module)?;
    types::type_builder::FieldType::define_in_ruby(&module)?;

    types::client_registry::ClientRegistry::define_in_ruby(&module)?;
    types::media::Audio::define_in_ruby(&module)?;
    types::media::Image::define_in_ruby(&module)?;
    types::media::Pdf::define_in_ruby(&module)?;
    types::media::Video::define_in_ruby(&module)?;

    // Register the new log collector classes
    types::log_collector::define_all_in_ruby(&module)?;

    // everything below this is for our own testing purposes
    module.define_module_function(
        "roundtrip",
        function!(ruby_to_json::RubyToJson::roundtrip, 1),
    )?;
    module.define_module_function(
        "serialize",
        function!(ruby_to_json::RubyToJson::serialize, 4),
    )?;

    Ok(())
}
